package org.setamv.springcloud.ribbon.rule;

import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.AbstractLoadBalancerRule;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.Server;

import java.util.List;
import java.util.Objects;

/**
 * 扩展ribbon的负载均衡规则：每个服务器重复访问N次后再访问下一个服务器
 */
public class MyLoadBalanceRule extends AbstractLoadBalancerRule {

    /**
     * 定义每个服务器重复访问的次数
     */
    public static final int REPEAT_TIMES = 3;

    private Server lastServer = null;
    private int visitTimes = 0;

    public MyLoadBalanceRule() {

    }

    public MyLoadBalanceRule(ILoadBalancer lb) {
        this();
        setLoadBalancer(lb);
    }

    @Override
    public Server choose(Object o) {
        List<Server> allServers = getLoadBalancer().getAllServers();
        int lastIndex = -1;
        if (lastServer != null) {
            int len = allServers.size();
            for (int i = 0; i < len; i++) {
                Server server = allServers.get(i);
                if (Objects.equals(server, lastServer)) {
                    lastIndex = i;
                    break;
                }
            }
            if (lastIndex != -1 && visitTimes < REPEAT_TIMES) {
                visitTimes++;
                return lastServer;
            }

            if (lastIndex >= len - 1) {
                lastIndex = -1;
            }
        }
        int nextIndex = lastIndex + 1;
        Server nextServer = allServers.get(nextIndex);
        visitTimes = 1;
        lastServer = nextServer;
        return nextServer;
    }

    @Override
    public void initWithNiwsConfig(IClientConfig clientConfig) {
    }
}
